import os
import argparse
import json
import cv2
from PIL import Image
from pylab import *
from utils.utils import get_yolo_boxes, makedirs
from utils.bbox import draw_boxes, get_boxes
from keras.models import load_model
# from tqdm import tqdm
import numpy as np
import time

# add by suyongsheng
import tensorflow as tf
from multiprocessing import cpu_count


class Yolo:
    def __init__(self, config_path):
        tf.Session(config=tf.ConfigProto(device_count={"CPU": cpu_count()}, inter_op_parallelism_threads=1,
                                         intra_op_parallelism_threads=2))
        with open(config_path) as config_buffer:
            self.config = json.load(config_buffer)

        self.net_h, self.net_w = 416, 416
        self.obj_thresh, self.nms_thresh = 0.5, 0.45
        os.environ['CUDA_VISIBLE_DEVICES'] = self.config['train']['gpus']
        self.infer_model = load_model(self.config['train']['saved_weights_name'])

    def predict(self, image, from_file=False, show_im=False):
        if from_file:
            image = cv2.imread(image)
        time_start = time.time()
        # predict the bounding boxes
        boxes = \
            get_yolo_boxes(self.infer_model, [image], self.net_h, self.net_w, self.config['model']['anchors'],
                           self.obj_thresh, self.nms_thresh)[0]
        time_end = time.time()
        print('yolo predict cost time=', time_end - time_start)
        image = draw_boxes(image, boxes, self.config['model']['labels'], self.obj_thresh)
        xyc_pack = get_boxes(image, boxes, self.config['model']['labels'], self.obj_thresh)
        if show_im:
            imshow(image)
            show()
        return xyc_pack

model_yolo = Yolo("./config.json")
image = cv2.imread("img00001.jpg")
model_yolo.predict(image,show_im=False)
model_yolo.predict(image,show_im=False)
model_yolo.predict(image,show_im=False)
model_yolo.predict(image,show_im=False)
model_yolo.predict(image,show_im=True)